# Copyright 2025 Huawei Technologies Co., Ltd
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Generate data for mcore MOE UT of inference."""
import numpy as np

def get_init_params(seq_len, batch_size, hidden_size, num_experts, moe_intermediate_size):
    """Generate initialization parameters"""
    np.random.seed(2025)
    gate_weight_shape = (num_experts, hidden_size)
    gate_e_score_correction_bias_shape = (num_experts,)
    experts_fc1_weight_shape = (num_experts, hidden_size, moe_intermediate_size * 2)
    experts_fc2_weight_shape = (num_experts, moe_intermediate_size, hidden_size)
    shared_experts_fc1_weight_shape = (moe_intermediate_size * 2, hidden_size)
    shared_experts_fc2_weight_shape = (hidden_size, moe_intermediate_size)
    return {
        "input": np.random.rand(batch_size * seq_len, hidden_size),
        "router.weight": 0.01 * np.random.rand(*gate_weight_shape),
        "router.expert_bias": 0.01 * np.random.rand(*gate_e_score_correction_bias_shape),
        "experts.weight1": 0.01 * np.random.rand(*experts_fc1_weight_shape),
        "experts.weight2": 0.01 * np.random.rand(*experts_fc2_weight_shape),
        "shared_experts.linear_fc1.weight": 0.01 * np.random.rand(*shared_experts_fc1_weight_shape),
        "shared_experts.linear_fc2.weight": 0.01 * np.random.rand(*shared_experts_fc2_weight_shape),
    }

def get_golden() -> dict[str, np.ndarray]:
    """Generate golden data for test."""
    tp1 = np.array(
        [[0.0004406, 0.000576, 0.000515, 0.0006485, 0.0004463, 0.000519,
          0.0006485, 0.000675, 0.00058, 0.0005875, 0.0005035, 0.0005417,
          0.000576, 0.0005035, 0.000572, 0.0005264, 0.0004463, 0.0004883,
          0.000538, 0.000511, 0.0005074, 0.0005417, 0.0005646, 0.000393,
          0.000599, 0.0004215, 0.000515, 0.0004272, 0.0004196, 0.0005417,
          0.0005493, 0.0004883],
         [0.000492, 0.0006523, 0.0005913, 0.00074, 0.000496, 0.0005875,
          0.000744, 0.0007706, 0.000641, 0.0006676, 0.0005684, 0.000618,
          0.000656, 0.000572, 0.0006523, 0.000595, 0.000515, 0.000557,
          0.0006104, 0.00058, 0.00058, 0.0006065, 0.000641, 0.0004444,
          0.000675, 0.0004578, 0.0005913, 0.0004864, 0.0004692, 0.0006104,
          0.0006256, 0.0005493],
         [0.0002708, 0.0003586, 0.0003242, 0.00041, 0.0002766, 0.0003242,
          0.0004044, 0.0004272, 0.0003586, 0.0003662, 0.000311, 0.0003338,
          0.000368, 0.0003128, 0.0003605, 0.00033, 0.0002823, 0.0003052,
          0.00033, 0.0003128, 0.0003242, 0.00033, 0.000351, 0.0002346,
          0.000368, 0.0002594, 0.0003262, 0.000269, 0.0002613, 0.0003433,
          0.0003433, 0.000307],
         [0.000288, 0.0003853, 0.0003376, 0.000431, 0.0002975, 0.0003471,
          0.0004234, 0.0004463, 0.0003872, 0.000389, 0.0003357, 0.0003548,
          0.0003777, 0.00033, 0.0003853, 0.0003471, 0.0002937, 0.0003185,
          0.0003548, 0.0003395, 0.0003376, 0.0003548, 0.0003738, 0.0002632,
          0.0003948, 0.0002766, 0.0003357, 0.0002785, 0.0002785, 0.0003586,
          0.000368, 0.0003204],
         [0.000328, 0.0004406, 0.0003967, 0.0005074, 0.0003433, 0.0004082,
          0.000492, 0.000519, 0.0004425, 0.000452, 0.0003815, 0.000414,
          0.0004501, 0.000391, 0.0004501, 0.0004196, 0.0003567, 0.0003872,
          0.000412, 0.0004005, 0.0003948, 0.0004177, 0.000431, 0.000307,
          0.000454, 0.0003204, 0.000389, 0.0003319, 0.0003185, 0.0004234,
          0.0004349, 0.0003757],
         [0.0005074, 0.0006638, 0.0005836, 0.0007477, 0.000511, 0.0006027,
          0.00074, 0.0007668, 0.0006676, 0.000683, 0.00058, 0.000618,
          0.00066, 0.00058, 0.0006676, 0.000614, 0.000511, 0.0005646,
          0.000614, 0.0005836, 0.0005836, 0.000614, 0.000641, 0.0004578,
          0.0006905, 0.000473, 0.0005836, 0.000496, 0.0004768, 0.0006332,
          0.0006332, 0.0005646],
         [0.0004272, 0.000557, 0.000492, 0.0006256, 0.0004349, 0.0005035,
          0.0006294, 0.0006523, 0.0005608, 0.000572, 0.0004845, 0.0005226,
          0.000557, 0.0004883, 0.0005646, 0.000515, 0.0004292, 0.0004807,
          0.000519, 0.000492, 0.0004883, 0.000511, 0.0005417, 0.0003853,
          0.0005875, 0.0003967, 0.0004997, 0.0004196, 0.0003967, 0.0005264,
          0.000534, 0.0004635],
         [0.000534, 0.0006905, 0.000622, 0.000786, 0.000534, 0.000622,
          0.000778, 0.0008125, 0.0006943, 0.0007095, 0.0006065, 0.0006447,
          0.000702, 0.0006027, 0.0006866, 0.0006332, 0.000538, 0.0005836,
          0.0006447, 0.0006027, 0.000614, 0.000637, 0.000679, 0.000454,
          0.000717, 0.0004864, 0.000622, 0.000519, 0.000496, 0.0006523,
          0.0006523, 0.0005875],
         [0.0004463, 0.0005875, 0.0005264, 0.0006638, 0.000454, 0.0005302,
          0.0006638, 0.0006943, 0.0005875, 0.0006027, 0.000511, 0.0005493,
          0.0005913, 0.000511, 0.0005875, 0.0005417, 0.0004597, 0.0005035,
          0.0005417, 0.0005226, 0.0005226, 0.0005417, 0.000572, 0.0003986,
          0.0006104, 0.0004158, 0.000534, 0.0004387, 0.0004196, 0.000557,
          0.0005684, 0.0004883],
         [0.0003872, 0.0005074, 0.0004616, 0.000576, 0.000393, 0.0004597,
          0.000576, 0.0006027, 0.0005074, 0.0005302, 0.0004463, 0.0004768,
          0.0005035, 0.0004482, 0.0005035, 0.0004616, 0.000393, 0.0004272,
          0.0004692, 0.000454, 0.0004501, 0.0004768, 0.0005035, 0.0003433,
          0.0005264, 0.0003624, 0.0004578, 0.0003853, 0.0003662, 0.0004807,
          0.0004692, 0.0004215],
         [0.0004005, 0.0005264, 0.000454, 0.0005875, 0.000412, 0.0004768,
          0.000595, 0.0006104, 0.000534, 0.0005417, 0.0004597, 0.000492,
          0.000519, 0.0004635, 0.0005302, 0.000496, 0.0004082, 0.0004559,
          0.000492, 0.0004616, 0.0004559, 0.0004883, 0.0005035, 0.000368,
          0.0005493, 0.0003796, 0.0004616, 0.0003967, 0.00037, 0.0004997,
          0.0005074, 0.0004463],
         [0.0004425, 0.000572, 0.0005074, 0.0006485, 0.000452, 0.0005264,
          0.0006447, 0.0006638, 0.00058, 0.000599, 0.000496, 0.000534,
          0.000572, 0.0005035, 0.00058, 0.000534, 0.0004425, 0.0004883,
          0.000534, 0.000511, 0.0005074, 0.000534, 0.0005608, 0.0004005,
          0.000595, 0.0004063, 0.0005035, 0.000431, 0.0004082, 0.0005417,
          0.000553, 0.0004787],
         [0.0002518, 0.00033, 0.0002918, 0.0003777, 0.0002594, 0.0003052,
          0.0003738, 0.0003853, 0.0003395, 0.0003471, 0.0002918, 0.000309,
          0.0003319, 0.0002937, 0.0003395, 0.0003166, 0.0002632, 0.0002918,
          0.000311, 0.0002995, 0.0002937, 0.000311, 0.0003223, 0.0002356,
          0.0003471, 0.0002375, 0.0002918, 0.0002518, 0.0002384, 0.0003204,
          0.0003223, 0.0002842],
         [0.0003662, 0.0004654, 0.0004272, 0.0005264, 0.0003624, 0.0004272,
          0.0005417, 0.0005646, 0.0004692, 0.0004883, 0.000412, 0.0004425,
          0.000475, 0.000414, 0.000471, 0.0004425, 0.00037, 0.0004158,
          0.0004406, 0.0004158, 0.0004158, 0.0004406, 0.0004616, 0.0003204,
          0.0004883, 0.0003395, 0.0004272, 0.0003567, 0.0003433, 0.000454,
          0.0004482, 0.0003967],
         [0.0002785, 0.000368, 0.00033, 0.0004196, 0.0002842, 0.0003338,
          0.0004158, 0.0004272, 0.0003662, 0.0003738, 0.0003166, 0.0003433,
          0.0003757, 0.0003262, 0.0003777, 0.0003433, 0.0002937, 0.0003185,
          0.0003433, 0.000328, 0.00033, 0.0003471, 0.0003529, 0.0002556,
          0.0003853, 0.0002594, 0.0003319, 0.0002766, 0.0002651, 0.000349,
          0.0003662, 0.0003147],
         [0.000389, 0.0005226, 0.0004635, 0.0005913, 0.000393, 0.0004673,
          0.00058, 0.0006104, 0.000511, 0.0005264, 0.000454, 0.0004883,
          0.000519, 0.000454, 0.000511, 0.0004692, 0.0004082, 0.0004387,
          0.000475, 0.0004578, 0.0004616, 0.0004807, 0.000511, 0.0003471,
          0.000534, 0.0003662, 0.0004692, 0.0003815, 0.00037, 0.0004864,
          0.000496, 0.0004349]], dtype=np.float16)

    tp1_no_shared = np.array(
        [[0.0003757, 0.0004425, 0.0004005, 0.000433, 0.0003433, 0.000393,
          0.0004387, 0.000473, 0.0004349, 0.0004501, 0.000351, 0.000393,
          0.0004387, 0.0003567, 0.0003872, 0.0004501, 0.0002937, 0.0003567,
          0.000393, 0.000311, 0.000351, 0.0003605, 0.0004215, 0.000286,
          0.0004272, 0.0003242, 0.0003185, 0.000288, 0.0003128, 0.0004177,
          0.0004063, 0.0003548],
         [0.0004158, 0.000492, 0.000452, 0.0004845, 0.0003757, 0.0004349,
          0.0004883, 0.0005226, 0.0004768, 0.0005035, 0.0003796, 0.0004406,
          0.000496, 0.000393, 0.0004368, 0.0005035, 0.000328, 0.000393,
          0.0004349, 0.0003471, 0.0003967, 0.000389, 0.000471, 0.0003223,
          0.0004692, 0.0003433, 0.0003586, 0.0003204, 0.0003433, 0.0004616,
          0.0004635, 0.000389],
         [0.0002317, 0.0002785, 0.0002556, 0.0002766, 0.0002117, 0.0002441,
          0.0002728, 0.0002995, 0.000269, 0.0002766, 0.0002165, 0.0002422,
          0.0002804, 0.0002213, 0.0002432, 0.0002785, 0.000185, 0.0002193,
          0.0002384, 0.0001907, 0.0002279, 0.0002184, 0.0002632, 0.0001707,
          0.0002613, 0.0001984, 0.000205, 0.0001793, 0.0001945, 0.0002651,
          0.0002575, 0.0002232],
         [0.000248, 0.0002956, 0.0002632, 0.00029, 0.0002317, 0.0002632,
          0.000288, 0.0003147, 0.0002937, 0.0002995, 0.0002346, 0.0002575,
          0.00029, 0.0002346, 0.0002632, 0.0002956, 0.0001945, 0.0002346,
          0.0002594, 0.0002098, 0.0002365, 0.0002365, 0.0002785, 0.0001917,
          0.0002823, 0.0002117, 0.0002089, 0.0001907, 0.000207, 0.0002785,
          0.0002728, 0.0002317],
         [0.0002804, 0.0003414, 0.000311, 0.0003452, 0.000267, 0.000311,
          0.0003357, 0.0003643, 0.0003338, 0.0003452, 0.000267, 0.0003052,
          0.0003433, 0.0002804, 0.000309, 0.0003567, 0.0002375, 0.0002842,
          0.0003052, 0.0002499, 0.0002747, 0.0002804, 0.0003262, 0.0002279,
          0.0003242, 0.0002441, 0.0002432, 0.0002251, 0.0002394, 0.000328,
          0.00033, 0.0002747],
         [0.000433, 0.0005074, 0.000454, 0.0004997, 0.0003948, 0.000454,
          0.0005035, 0.000534, 0.0005035, 0.0005226, 0.0004044, 0.0004501,
          0.0004997, 0.000412, 0.000454, 0.0005226, 0.0003357, 0.00041,
          0.0004501, 0.0003567, 0.0004025, 0.0004044, 0.0004807, 0.0003376,
          0.000492, 0.0003605, 0.0003605, 0.0003357, 0.0003548, 0.0004883,
          0.0004692, 0.0004082],
         [0.0003643, 0.0004292, 0.0003853, 0.0004177, 0.0003319, 0.0003777,
          0.0004253, 0.000452, 0.0004196, 0.0004387, 0.0003395, 0.0003815,
          0.0004215, 0.0003471, 0.0003815, 0.0004368, 0.0002804, 0.0003452,
          0.0003777, 0.0003014, 0.0003376, 0.0003395, 0.0004044, 0.0002842,
          0.0004177, 0.0003014, 0.000307, 0.0002823, 0.0002975, 0.0004044,
          0.0003986, 0.0003357],
         [0.000452, 0.0005302, 0.0004787, 0.0005226, 0.0004025, 0.0004635,
          0.000519, 0.0005646, 0.0005226, 0.000538, 0.0004158, 0.0004635,
          0.0005264, 0.0004177, 0.0004616, 0.000538, 0.0003471, 0.00041,
          0.0004616, 0.0003643, 0.0004234, 0.000412, 0.0005035, 0.0003319,
          0.0005035, 0.000372, 0.0003853, 0.0003433, 0.0003643, 0.0004997,
          0.0004845, 0.0004196],
         [0.0003796, 0.000452, 0.0004082, 0.0004425, 0.0003471, 0.0003986,
          0.0004463, 0.0004845, 0.0004406, 0.0004616, 0.0003567, 0.0003986,
          0.0004482, 0.0003586, 0.0003986, 0.0004616, 0.0002995, 0.0003605,
          0.000393, 0.0003185, 0.0003624, 0.0003548, 0.0004292, 0.00029,
          0.0004349, 0.0003147, 0.0003319, 0.0002956, 0.0003128, 0.0004253,
          0.0004253, 0.000351],
         [0.000328, 0.0003815, 0.0003529, 0.0003738, 0.0002956, 0.0003395,
          0.0003796, 0.000412, 0.0003796, 0.0003986, 0.0003033, 0.0003395,
          0.0003796, 0.000309, 0.0003376, 0.000389, 0.0002499, 0.0003033,
          0.0003357, 0.0002708, 0.000309, 0.000307, 0.00037, 0.0002499,
          0.00037, 0.0002708, 0.000286, 0.0002594, 0.0002728, 0.0003643,
          0.0003471, 0.0002975],
         [0.0003414, 0.00041, 0.0003586, 0.0003986, 0.0003185, 0.0003624,
          0.0004025, 0.0004215, 0.0004044, 0.000414, 0.0003262, 0.0003624,
          0.000393, 0.0003338, 0.0003605, 0.0004234, 0.000267, 0.000328,
          0.0003643, 0.000288, 0.0003147, 0.000328, 0.0003815, 0.0002747,
          0.000393, 0.00029, 0.0002804, 0.000267, 0.0002785, 0.0003853,
          0.0003796, 0.000328],
         [0.0003796, 0.0004444, 0.0003967, 0.0004349, 0.0003471, 0.0003986,
          0.0004406, 0.0004692, 0.0004425, 0.0004616, 0.0003529, 0.000389,
          0.0004368, 0.0003567, 0.0004005, 0.0004597, 0.0002975, 0.0003548,
          0.000393, 0.0003223, 0.0003548, 0.0003548, 0.0004215, 0.0003014,
          0.0004272, 0.0003128, 0.0003185, 0.0002956, 0.000309, 0.0004215,
          0.0004196, 0.0003471],
         [0.0002155, 0.0002556, 0.0002298, 0.0002537, 0.0001993, 0.0002308,
          0.0002537, 0.000267, 0.0002556, 0.0002651, 0.000206, 0.000227,
          0.0002518, 0.0002108, 0.0002308, 0.000269, 0.0001717, 0.0002117,
          0.0002289, 0.0001869, 0.0002022, 0.000207, 0.0002432, 0.0001764,
          0.000248, 0.0001793, 0.0001802, 0.0001707, 0.0001783, 0.000248,
          0.0002422, 0.0002079],
         [0.0003128, 0.0003643, 0.0003395, 0.0003586, 0.0002804, 0.0003262,
          0.00037, 0.0003986, 0.0003586, 0.0003796, 0.0002918, 0.000328,
          0.0003662, 0.0002975, 0.000328, 0.0003796, 0.000246, 0.0003033,
          0.0003262, 0.0002594, 0.0002956, 0.0002975, 0.000351, 0.0002403,
          0.0003548, 0.0002632, 0.0002747, 0.0002441, 0.0002632, 0.0003529,
          0.0003433, 0.0002937],
         [0.0002346, 0.0002804, 0.0002556, 0.0002766, 0.0002146, 0.000248,
          0.0002766, 0.0002918, 0.0002708, 0.0002823, 0.0002165, 0.000246,
          0.0002804, 0.0002279, 0.0002518, 0.00029, 0.0001888, 0.0002241,
          0.000248, 0.0001993, 0.0002222, 0.0002241, 0.0002613, 0.000186,
          0.000267, 0.0001955, 0.0001984, 0.0001822, 0.0001936, 0.0002651,
          0.000269, 0.0002251],
         [0.000328, 0.000393, 0.0003529, 0.000389, 0.0002995, 0.0003452,
          0.0003834, 0.0004196, 0.0003815, 0.0003986, 0.0003052, 0.000351,
          0.000389, 0.0003128, 0.0003376, 0.0003967, 0.0002613, 0.000311,
          0.0003414, 0.0002708, 0.0003128, 0.000307, 0.0003796, 0.000248,
          0.000372, 0.0002728, 0.000286, 0.0002537, 0.0002708, 0.00037,
          0.0003624, 0.000309]], dtype=np.float16)

    return {
        "tp1": tp1,
        "tp1_no_shared": tp1_no_shared,
    }

GOLDEN_DATA = get_golden()
